#!/usr/bin/env python3
import os
import cv2
import numpy as np
from PIL import Image
import time
import glob
import queue
import threading
from display import PositionVisualizer
from display_top import PoseVisualizer
from shared_memory_reader import SharedMemoryReader
from seg_pose import SegPose
from compute_position import ComputePosition
from utils import get_depth_image
from logger import setup_logger

logger = setup_logger()

class ImageProcessor:
    def __init__(self, from_folder=True, color_folder="./data/Color1", depth_folder="./data/Depth1"):
        self.from_folder = from_folder
        self.color_folder = color_folder
        self.depth_folder = depth_folder
        self.seg_pose = SegPose()
        self.compute_position = ComputePosition()
        self.position_visualizer = PositionVisualizer()
        self.pose_visualizer = PoseVisualizer()
        
        # 队列容量增大以缓冲更多图像
        self.image_queue = queue.Queue(maxsize=200)
        self.processed_queue = queue.Queue(maxsize=200)
        self.position_queue = queue.Queue(maxsize=200)
        
        # 统计信息
        self.frame_count = 0
        self.saved_count = 0
        self.start_time = time.time()
        self.running = True
        
    def set_standard_position(self):
        """设置标准位置"""
        st_img = cv2.imread("./data/Color1/440.png")
        st_depth_img = get_depth_image("./data/Color1/440.png")
        
        logger.info("==========设置标准位置==========")
        seg_result, pose_result = self.seg_pose.get_result(st_img)
        if seg_result is None or pose_result is None:
            logger.error("==========标准图像检测失败==========")
            return None, None
            
        st_pose, st_vector = self.compute_position.set_standard(st_img, st_depth_img, seg_result, pose_result)
        if st_pose is None or st_vector is None:
            logger.error("==========标准位置设定失败==========")
        else:
            logger.info("==========标准位置设定成功==========")
            logger.info(f"标准位置: {st_pose}")
            logger.info(f"标准向量: {st_vector}")
        return st_pose, st_vector

    def image_reader(self):
        """图像读取线程（从文件夹或共享内存）"""
        if self.from_folder:
            color_files = sorted(glob.glob(os.path.join(self.color_folder, "*.png")))
            for color_path in color_files:
                if not self.running:
                    break
                    
                depth_path = os.path.join(self.depth_folder, os.path.basename(color_path))
                if not os.path.exists(depth_path):
                    logger.warning(f"缺失深度图: {depth_path}")
                    continue
                    
                color_img = cv2.imread(color_path)
                depth_img = get_depth_image(color_path)
                timestamp = os.path.splitext(os.path.basename(color_path))[0]
                
                if color_img is not None and depth_img is not None:
                    self.image_queue.put((timestamp, color_path, color_img, depth_img))
                else:
                    logger.warning(f"图像读取失败: {color_path} 或 {depth_path}")
        else:
            shm_reader = SharedMemoryReader()
            try:
                while self.running:
                    readable_timestamp, timestamp, color_img, depth_img = shm_reader.read_images()
                    if color_img is not None and depth_img is not None:
                        success, color_path, depth_path = shm_reader.save_images(timestamp, color_img, depth_img)
                        if success:
                            self.saved_count += 1
                        self.image_queue.put((readable_timestamp, color_path, color_img, depth_img))
                    else:
                        time.sleep(0.01)
            finally:
                shm_reader.close()

    def image_processor(self):
        """图像处理线程（分割和姿态估计）"""
        while self.running:
            try:
                timestamp, color_path, color_img, depth_img = self.image_queue.get(timeout=1.0)
                self.frame_count += 1
                
                step_start = time.time()
                logger.info(f"[{self.frame_count}] 处理图像: {timestamp}, 彩色: {color_img.shape}, 深度: {depth_img.shape}")
                
                t1 = time.time()
                seg_result, pose_result = self.seg_pose.get_result(color_img)
                t2 = time.time()
                logger.info(f"[{self.frame_count}] 分割+姿态估计耗时: {(t2 - t1) * 1000:.2f} ms")
                
                if seg_result is None or pose_result is None:
                    logger.warning(f"[{self.frame_count}] 分割或姿态估计失败: {timestamp}")
                    continue
                
                self.processed_queue.put((timestamp, color_path, color_img, depth_img, seg_result, pose_result))
                logger.info(f"[{self.frame_count}] 处理总耗时: {(time.time() - step_start) * 1000:.2f} ms")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"[{self.frame_count}] 处理错误: {e}")

    def position_computer(self, st_pose, st_vector):
        """位置计算线程"""
        while self.running:
            try:
                timestamp, color_path, color_img, depth_img, seg_result, pose_result = self.processed_queue.get(timeout=1.0)
                
                t3 = time.time()
                pose, vector = self.compute_position.process_image_pair(timestamp, color_img, depth_img, seg_result, pose_result)
                t4 = time.time()
                logger.info(f"[{self.frame_count}] 位置计算耗时: {(t4 - t3) * 1000:.2f} ms")
                
                logger.info(f"位置: {pose}")
                logger.info(f"向量: {vector}")
                
                self.position_queue.put((timestamp, color_path, pose, vector, st_pose, st_vector))
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"[{self.frame_count}] 位置计算错误: {e}")

    def visualizer(self):
        """可视化线程"""
        while self.running:
            try:
                timestamp, color_path, pose, vector, st_pose, st_vector = self.position_queue.get(timeout=1.0)
                
                t5 = time.time()
                self.pose_visualizer.visualize_pose(pose, st_pose, color_path)
                self.position_visualizer.visualize_position(st_vector, vector, color_path)
                t6 = time.time()
                logger.info(f"[{self.frame_count}] 可视化耗时: {(t6 - t5) * 1000:.2f} ms")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"[{self.frame_count}] 可视化错误: {e}")

    def run(self):
        """启动所有线程"""
        st_pose, st_vector = self.set_standard_position()
        if st_pose is None or st_vector is None:
            return
            
        threads = [
            threading.Thread(target=self.image_reader, daemon=True),
            threading.Thread(target=self.image_processor, daemon=True),
            # 可选：添加更多 image_processor 线程（需测试 NPU 并发性能）
            # threading.Thread(target=self.image_processor, daemon=True),
            threading.Thread(target=self.position_computer, args=(st_pose, st_vector), daemon=True),
            threading.Thread(target=self.visualizer, daemon=True)
        ]
        
        for thread in threads:
            thread.start()
            
        try:
            while True:
                time.sleep(1)
                if self.frame_count > 0:
                    elapsed_time = time.time() - self.start_time
                    fps = self.frame_count / elapsed_time
                    logger.info(f"=== 统计信息 ===")
                    logger.info(f"处理帧数: {self.frame_count}, 保存帧数: {self.saved_count}")
                    logger.info(f"平均帧率: {fps:.1f}, 保存成功率: {self.saved_count/self.frame_count*100:.1f}%")
                    logger.info(f"=================")
                
        except KeyboardInterrupt:
            logger.info("收到中断信号，正在关闭...")
            self.running = False
            
        finally:
            self.seg_pose.shutdown()  # 关闭 SegPose
            for thread in threads:
                thread.join(timeout=2.0)
                
            total_time = time.time() - self.start_time
            if self.frame_count > 0:
                avg_fps = self.frame_count / total_time
                logger.info(f"==========程序结束==========")
                logger.info(f"总处理帧数: {self.frame_count}")
                logger.info(f"平均帧率: {avg_fps:.2f}")
            else:
                logger.info("未处理任何帧")

if __name__ == "__main__":
    processor = ImageProcessor(from_folder=True)
    processor.run()